#!/usr/bin/env python3
"""
数据库连接池管理模块

本模块负责管理 MySQL 数据库连接池，提供连接的创建、验证、重用和释放功能。
主要功能包括：
- 连接池管理
- 连接健康检查
- 自动重连机制
- 连接超时处理
- 性能监控
"""

import time
import threading
from typing import Optional, Dict, Any, List
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timedelta

import mysql.connector
from mysql.connector import Error as MySQLError
from mysql.connector.pooling import MySQLConnectionPool

from .config import DatabaseConfig
from .logger import get_logger, log_database_connection, log_database_connection_success, log_database_connection_error


class ConnectionPoolError(Exception):
    """
    连接池错误异常
    
    当连接池操作失败时抛出此异常。
    """
    
    def __init__(self, message: str, original_error: Optional[Exception] = None):
        """
        初始化连接池错误
        
        Args:
            message: 错误消息
            original_error: 原始异常
        """
        self.message = message
        self.original_error = original_error
        super().__init__(message)


class QueryExecutionError(Exception):
    """
    查询执行错误异常
    
    当查询执行失败时抛出此异常。
    """
    
    def __init__(self, message: str, query: Optional[str] = None, original_error: Optional[Exception] = None):
        """
        初始化查询执行错误
        
        Args:
            message: 错误消息
            query: 导致错误的查询语句
            original_error: 原始异常
        """
        self.message = message
        self.query = query
        self.original_error = original_error
        super().__init__(message)


@dataclass
class ConnectionStats:
    """
    连接统计信息类
    
    记录连接池的使用统计信息，包括连接数、执行次数、错误次数等。
    """
    
    total_connections: int = field(default=0)
    active_connections: int = field(default=0)
    successful_queries: int = field(default=0)
    failed_queries: int = field(default=0)
    total_execution_time: float = field(default=0.0)
    last_error: Optional[str] = field(default=None)
    last_error_time: Optional[datetime] = field(default=None)
    created_at: datetime = field(default_factory=datetime.now)
    
    def add_query_success(self, execution_time: float):
        """
        记录成功查询
        
        Args:
            execution_time: 查询执行时间（秒）
        """
        self.successful_queries += 1
        self.total_execution_time += execution_time
    
    def add_query_failure(self, error_message: str):
        """
        记录失败查询
        
        Args:
            error_message: 错误信息
        """
        self.failed_queries += 1
        self.last_error = error_message
        self.last_error_time = datetime.now()
    
    def get_average_execution_time(self) -> float:
        """
        获取平均执行时间
        
        Returns:
            float: 平均执行时间（秒）
        """
        if self.successful_queries == 0:
            return 0.0
        return self.total_execution_time / self.successful_queries
    
    def get_success_rate(self) -> float:
        """
        获取成功率
        
        Returns:
            float: 成功率（0-1之间）
        """
        total_queries = self.successful_queries + self.failed_queries
        if total_queries == 0:
            return 1.0
        return self.successful_queries / total_queries


class ConnectionPool:
    """
    数据库连接池管理类
    
    管理 MySQL 数据库连接池，提供连接的获取、释放、健康检查等功能。
    支持自动重连、连接验证、性能监控等特性。
    """
    
    def __init__(self, config: DatabaseConfig):
        """
        初始化连接池
        
        Args:
            config: 数据库配置实例
        """
        self.config = config
        self.logger = get_logger(__name__)
        self.stats = ConnectionStats()
        self._pool: Optional[MySQLConnectionPool] = None
        self._lock = threading.RLock()
        self._last_health_check = datetime.now()
        self._health_check_interval = timedelta(minutes=5)
        
        # 初始化连接池
        self._initialize_pool()
    
    def _initialize_pool(self):
        """
        初始化数据库连接池
        
        创建 MySQL 连接池实例，配置连接参数和池大小。
        
        Raises:
            MySQLError: 当连接池创建失败时抛出异常
        """
        try:
            connection_params = self.config.to_connection_params()
            
            # 创建连接池配置
            pool_config = {
                **connection_params,
                'pool_name': 'mysql_mcp_pool',
                'pool_size': self.config.max_connections,
                'pool_reset_session': True,
                'autocommit': True,
            }
            
            self.logger.info(
                f"正在创建连接池: {self.config.host}:{self.config.port}/{self.config.database}"
            )
            
            self._pool = MySQLConnectionPool(**pool_config)
            self.stats.total_connections = self.config.max_connections
            
            # 测试连接
            self._test_connection()
            
            self.logger.info(
                f"连接池创建成功，最大连接数: {self.config.max_connections}"
            )
            
        except MySQLError as e:
            error_msg = f"连接池初始化失败: {str(e)}"
            self.logger.error(error_msg)
            self.stats.add_query_failure(error_msg)
            raise MySQLError(error_msg) from e
    
    def _test_connection(self):
        """
        测试数据库连接
        
        通过执行简单查询来验证数据库连接是否正常。
        
        Raises:
            MySQLError: 当连接测试失败时抛出异常
        """
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                cursor.close()
                
                if result != (1,):
                    raise MySQLError("连接测试失败：查询结果异常")
                
                self.logger.debug("数据库连接测试成功")
                
        except MySQLError as e:
            error_msg = f"数据库连接测试失败: {str(e)}"
            self.logger.error(error_msg)
            raise MySQLError(error_msg) from e
    
    @contextmanager
    def get_connection(self):
        """
        获取数据库连接（上下文管理器）
        
        使用上下文管理器自动管理连接的获取和释放。
        支持重试机制和连接健康检查。
        
        Yields:
            mysql.connector.MySQLConnection: 数据库连接对象
            
        Raises:
            MySQLError: 当无法获取有效连接时抛出异常
        """
        connection = None
        retry_count = 0
        
        while retry_count <= self.config.retry_count:
            try:
                with self._lock:
                    # 检查连接池健康状态
                    self._check_pool_health()
                    
                    # 获取连接
                    connection = self._pool.get_connection()
                    self.stats.active_connections += 1
                
                # 验证连接
                if not connection.is_connected():
                    connection.reconnect(attempts=1, delay=0)
                
                self.logger.debug(f"获取数据库连接成功，活跃连接数: {self.stats.active_connections}")
                
                try:
                    yield connection
                finally:
                    # 确保连接被正确释放
                    if connection and connection.is_connected():
                        connection.close()
                    
                    with self._lock:
                        self.stats.active_connections = max(0, self.stats.active_connections - 1)
                
                return  # 成功获取并使用连接，退出重试循环
                
            except MySQLError as e:
                retry_count += 1
                error_msg = f"获取数据库连接失败 (尝试 {retry_count}/{self.config.retry_count + 1}): {str(e)}"
                
                if connection:
                    try:
                        connection.close()
                    except:
                        pass
                    
                    with self._lock:
                        self.stats.active_connections = max(0, self.stats.active_connections - 1)
                
                if retry_count <= self.config.retry_count:
                    self.logger.warning(f"{error_msg}，{self.config.retry_delay} 秒后重试")
                    time.sleep(self.config.retry_delay)
                else:
                    self.logger.error(error_msg)
                    self.stats.add_query_failure(error_msg)
                    raise MySQLError(error_msg) from e
    
    def _check_pool_health(self):
        """
        检查连接池健康状态
        
        定期检查连接池的健康状态，必要时重新初始化连接池。
        """
        now = datetime.now()
        
        # 检查是否需要进行健康检查
        if now - self._last_health_check < self._health_check_interval:
            return
        
        self._last_health_check = now
        
        try:
            # 检查连接池是否存在
            if not self._pool:
                self.logger.warning("连接池不存在，正在重新初始化")
                self._initialize_pool()
                return
            
            # 简单的健康检查
            self.logger.debug("执行连接池健康检查")
            
        except Exception as e:
            self.logger.error(f"连接池健康检查失败: {str(e)}")
            # 尝试重新初始化连接池
            try:
                self._initialize_pool()
            except Exception as init_error:
                self.logger.error(f"连接池重新初始化失败: {str(init_error)}")
    
    def execute_query(self, query: str, params: Optional[tuple] = None) -> List[Dict[str, Any]]:
        """
        执行 SQL 查询
        
        Args:
            query: SQL 查询语句
            params: 查询参数
            
        Returns:
            List[Dict[str, Any]]: 查询结果列表
            
        Raises:
            MySQLError: 当查询执行失败时抛出异常
        """
        start_time = time.time()
        
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor(dictionary=True)
                
                if params:
                    cursor.execute(query, params)
                else:
                    cursor.execute(query)
                
                # 获取查询结果
                if cursor.with_rows:
                    results = cursor.fetchall()
                else:
                    results = []
                
                cursor.close()
                
                execution_time = time.time() - start_time
                self.stats.add_query_success(execution_time)
                
                self.logger.debug(
                    f"查询执行成功，耗时: {execution_time:.3f}秒，结果数: {len(results)}"
                )
                
                return results
                
        except MySQLError as e:
            execution_time = time.time() - start_time
            error_msg = f"查询执行失败: {str(e)}"
            
            self.logger.error(f"{error_msg}，耗时: {execution_time:.3f}秒")
            self.stats.add_query_failure(error_msg)
            
            raise MySQLError(error_msg) from e
    
    def get_table_list(self) -> List[str]:
        """
        获取数据库表列表
        
        Returns:
            List[str]: 表名列表
            
        Raises:
            MySQLError: 当获取表列表失败时抛出异常
        """
        try:
            query = "SHOW TABLES"
            results = self.execute_query(query)
            
            # 提取表名
            table_key = f"Tables_in_{self.config.database}"
            tables = [row[table_key] for row in results if table_key in row]
            
            self.logger.debug(f"获取到 {len(tables)} 个数据表")
            return tables
            
        except MySQLError as e:
            error_msg = f"获取表列表失败: {str(e)}"
            self.logger.error(error_msg)
            raise MySQLError(error_msg) from e
    
    def get_table_schema(self, table_name: str) -> List[Dict[str, Any]]:
        """
        获取表结构信息
        
        Args:
            table_name: 表名
            
        Returns:
            List[Dict[str, Any]]: 表结构信息列表
            
        Raises:
            MySQLError: 当获取表结构失败时抛出异常
        """
        try:
            query = "DESCRIBE %s" % table_name  # 注意：这里应该使用参数化查询
            results = self.execute_query(f"DESCRIBE `{table_name}`")
            
            self.logger.debug(f"获取表 {table_name} 结构信息成功")
            return results
            
        except MySQLError as e:
            error_msg = f"获取表 {table_name} 结构失败: {str(e)}"
            self.logger.error(error_msg)
            raise MySQLError(error_msg) from e
    
    def execute_batch_queries(self, queries: List[tuple]) -> List[Dict[str, Any]]:
        """
        批量执行 SQL 查询
        
        Args:
            queries: 查询列表，每个元素为 (query, params) 元组
            
        Returns:
            List[Dict[str, Any]]: 批量查询结果列表
            
        Raises:
            MySQLError: 当批量查询执行失败时抛出异常
        """
        start_time = time.time()
        results = []
        
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor(dictionary=True)
                
                for query, params in queries:
                    try:
                        if params:
                            cursor.execute(query, params)
                        else:
                            cursor.execute(query)
                        
                        # 获取查询结果
                        if cursor.with_rows:
                            result = cursor.fetchall()
                        else:
                            result = {
                                'affected_rows': cursor.rowcount,
                                'last_insert_id': cursor.lastrowid
                            }
                        
                        results.append(result)
                        
                    except MySQLError as e:
                        # 记录单个查询失败，但继续执行其他查询
                        error_result = {
                            'error': str(e),
                            'query': query,
                            'params': params
                        }
                        results.append(error_result)
                        self.logger.warning(f"批量查询中的单个查询失败: {str(e)}")
                
                cursor.close()
                
                execution_time = time.time() - start_time
                self.stats.add_query_success(execution_time)
                
                self.logger.debug(
                    f"批量查询执行完成，耗时: {execution_time:.3f}秒，查询数: {len(queries)}"
                )
                
                return results
                
        except MySQLError as e:
            execution_time = time.time() - start_time
            error_msg = f"批量查询执行失败: {str(e)}"
            
            self.logger.error(f"{error_msg}，耗时: {execution_time:.3f}秒")
            self.stats.add_query_failure(error_msg)
            
            raise MySQLError(error_msg) from e
    
    def get_connection_info(self) -> Dict[str, Any]:
        """
        获取数据库连接信息
        
        Returns:
            Dict[str, Any]: 包含数据库连接信息的字典
        """
        return {
            'host': self.config.host,
            'port': self.config.port,
            'database': self.config.database,
            'user': self.config.user,
            'charset': self.config.charset,
            'pool_size': self.config.pool_size,
            'autocommit': self.config.autocommit
        }
    
    def get_stats(self) -> ConnectionStats:
        """
        获取连接池统计信息
        
        Returns:
            ConnectionStats: 统计信息实例
        """
        return self.stats
    
    async def initialize(self):
        """
        异步初始化连接池
        
        为了与MCP服务器兼容，提供异步初始化方法。
        实际上调用同步的_initialize_pool方法。
        """
        # 连接池在__init__中已经初始化，这里只是为了兼容性
        pass
    
    def close(self):
        """
        关闭连接池
        
        释放所有连接资源，清理连接池。
        """
        try:
            with self._lock:
                if self._pool:
                    # MySQL 连接池没有直接的 close 方法
                    # 我们只需要将引用设为 None，让垃圾回收器处理
                    self._pool = None
                    self.stats.active_connections = 0
                    
                    self.logger.info("连接池已关闭")
                    
        except Exception as e:
            self.logger.error(f"关闭连接池时发生错误: {str(e)}")
    
    def __enter__(self):
        """
        上下文管理器入口
        
        Returns:
            ConnectionPool: 连接池实例
        """
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        上下文管理器出口
        
        Args:
            exc_type: 异常类型
            exc_val: 异常值
            exc_tb: 异常追踪
        """
        self.close()


# 为了保持向后兼容性，创建别名
DatabaseConnectionPool = ConnectionPool